8.5 Диффузионные модели

Введение

В этом параграфе мы снова попробуем решить задачу генерации, когда нам дана выборка объектов из распределения x0q(x)x_0 \sim q(x), и хотим научиться генерировать новые объекты из распределения , которых нет в нашей выборке.

Вероятно, вы уже знакомы с другими генеративными моделями, например VAE или GAN-ы. Здесь же мы познакомим вас с еще одним видом генеративных моделей: диффузионные модели, которые стали крайне популярны в последнее время благодаря своему высокому качеству генерации объектов из заданного распределения. В общий чертах, они работают следующим образом: берем шум из N(0,I)\mathcal{N}(0, I) и шаг за шагом удаляем компоненты шума до тех пор, пока не получим объект x0x_0 из распределения, см. иллюстрацию ниже.

Screenshot

Более детально

Для детального понимания стоит объяснить, что такое прямой и обратный диффузионные процессы. Прямой процесс заключается в постепенном зашумлении картинки с помощью распределения qq, а обратный, наоборот, в расшумлении с помощью распределения pp. Их можно схематично изобразить следующим образом:

Artboard

Прямой диффузионный процесс определяется как апостериорное распределение q(x1:Tx0)q(x_{1:T}|x_0). Это распределение также является Марковской цепочкой, которая постепенно добавляет гауссовский шум к объекту x0x_0. На каждом шаге шум добавляется с различной магнитудой, которая определяется расписанием дисперсий {β1,...,βT}\{\beta_1, ... , \beta_T\}. При правильном выборе расписания в пределе по числу шагов TT мы должны сойтись к шуму из N(0,I)\mathcal{N}(0, I). В качестве распределений qq берут нормальные распределения:

q(xtxt1):=N(xt;1βtxt1,βtI),       q(x1:Tx0)=t=1Tq(xtxt1)q(x_t | x_{t - 1}) := \mathcal{N}(x_t; \sqrt{1 - \beta_t}x_{t - 1}, \beta_tI), \ \ \ \ \ \ \ q(x_{1:T}|x_0) = \prod_{t = 1}^T q(x_t | x_{t - 1})

Теперь перейдем к обратному процессу и к самой диффузионной модели.

Диффузионная модель - это вероятностная модель с латентными переменными вида pθ(x0):=pθ(x0:T)dx1:Tp_\theta(x_0) := \int p_\theta(x_{0:T}) dx_{1:T}, где промежуточные состояния x1,...,xTx_1, ..., x_T соответствуют зашумленным объектам, a x0x_0 - объект из распределения. Совместное распределение pθ(x0:T)p_\theta(x_{0:T}) называет обратным диффузионным процессом, который представляет собой Марковскую цепочку из гауссовских распределений pθ(xi1xi)p_\theta(x_{i-1}|x_{i}):

p(x0:T)=p(x0)t=1Tpθ(xt1xt)         pθ(xT)=N(xT0,I)p(x_{0:T}) = p(x_0) \prod_{t = 1}^Tp_{\theta}(x_{t-1}|x_t) \ \ \ \ \ \ \ \ \ p_\theta(x_{T})=\mathcal{N}(x_T | 0, I)

pθ(xt1xt):=N(xt1;μθ(xt,t),Σθ(xt,t)) p_{\theta}(x_{t - 1}|x_t):= \mathcal{N}(x_{t - 1}; \mu_{\theta}(x_t, t), \Sigma_{\theta}(x_t, t))

Таким образом, обратный процесс параметризуется моделью θ\theta, которая по зашумленному объекту xtx_t и шагу tt предсказывает среднее μθ(xt,t)\mu_{\theta}(x_t, t) и дисперсию Σθ(xt,t)\Sigma_{\theta}(x_t, t).

Обучение диффузионной модели

Диффузионный модели обучаются, максимизируя вариационную нижнюю оценку (ELBO) логарифма правдоподобия logpθ(x0)\log p_{\theta}(x_0). По тому же принципу обучаются VAE, с тем лишь отличием, что у диффузионных моделей другая форма модели с латентными переменными. Итак, давайте выведем ELBO для диффузии:

logpθ(x0)logpθ(x0)+DKL(q(x1:Tx0)pθ(x1:Tx0))=logpθ(x0)+Ex1:Tq(x1:Tx0)[logq(x1:Tx0)pθ(x0)pθ(x0:T)]=logpθ(x0)+Eq[logq(x1:Tx0)pθ(x0:T)+logpθ(x0)]=Eq[logq(x1:Tx0)pθ(x0:T)]Let LVLB=Eq(x0:T)[logq(x1:Tx0)pθ(x0:T)]Eq(x0)logpθ(x0)\begin{aligned} - \log p_\theta(\mathbf{x}_0) &\leq - \log p_\theta(\mathbf{x}_0) + D_\text{KL}(q(\mathbf{x}_{1:T}\vert\mathbf{x}_0) \| p_\theta(\mathbf{x}_{1:T}\vert\mathbf{x}_0) ) \\ &= -\log p_\theta(\mathbf{x}_0) + \mathbb{E}_{\mathbf{x}_{1:T}\sim q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} \Big[ \log\frac{q(\mathbf{x}_{1:T}\vert\mathbf{x}_0) p_\theta(\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})} \Big] \\ &= -\log p_\theta(\mathbf{x}_0) + \mathbb{E}_q \Big[ \log\frac{q(\mathbf{x}_{1:T}\vert\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})} + \log p_\theta(\mathbf{x}_0) \Big] \\ &= \mathbb{E}_q \Big[ \log \frac{q(\mathbf{x}_{1:T}\vert\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})} \Big] \\ \text{Let }L_\text{VLB} &= \mathbb{E}_{q(\mathbf{x}_{0:T})} \Big[ \log \frac{q(\mathbf{x}_{1:T}\vert\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})} \Big] \geq - \mathbb{E}_{q(\mathbf{x}_0)} \log p_\theta(\mathbf{x}_0) \end{aligned}

  • Комментарий

Если вы знакомы с VAE, то вывод LVLBL_{VLB} должен быть вам понятен, однако ниже приведен вывод с помощью неравенства Йенсена

LCE=Eq(x0)logpθ(x0)=Eq(x0)log(pθ(x0:T)dx1:T)=Eq(x0)log(q(x1:Tx0)pθ(x0:T)q(x1:Tx0)dx1:T)=Eq(x0)log(Eq(x1:Tx0)pθ(x0:T)q(x1:Tx0))Eq(x0:T)logpθ(x0:T)q(x1:Tx0)=Eq(x0:T)[logq(x1:Tx0)pθ(x0:T)]=LVLB\begin{aligned} L_\text{CE} &= - \mathbb{E}_{q(\mathbf{x}_0)} \log p_\theta(\mathbf{x}_0) \\ &= - \mathbb{E}_{q(\mathbf{x}_0)} \log \Big( \int p_\theta(\mathbf{x}_{0:T}) d\mathbf{x}_{1:T} \Big) \\ &= - \mathbb{E}_{q(\mathbf{x}_0)} \log \Big( \int q(\mathbf{x}_{1:T} \vert \mathbf{x}_0) \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_{0})} d\mathbf{x}_{1:T} \Big) \\ &= - \mathbb{E}_{q(\mathbf{x}_0)} \log \Big( \mathbb{E}_{q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_{0})} \Big) \\ &\leq - \mathbb{E}_{q(\mathbf{x}_{0:T})} \log \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_{0})} \\ &= \mathbb{E}_{q(\mathbf{x}_{0:T})}\Big[\log \frac{q(\mathbf{x}_{1:T} \vert \mathbf{x}_{0})}{p_\theta(\mathbf{x}_{0:T})} \Big] = L_\text{VLB} \end{aligned}

Теперь вернемся к распределению q(xtxt1)q(x_t | x_{t - 1}). Для того чтобы получить xtx_t, придется итеративно получать x1,...,xt1x_1, ..., x_{t - 1}. Однако это можно сделать более эффективно благодаря нормальным распределениям. Для этого обозначим αt:=1βt\alpha_t := 1- \beta_t и αˉt:=i=1tαi\bar{\alpha}_t:= \prod_{i = 1}^t\alpha_i, тогда

q(xtx0)=N(xt;αˉtx0,(1αˉt)I)q(x_t | x_0) = \mathcal{N}(x_t;\sqrt{\bar{\alpha}_t} x_0, (1-\bar{\alpha}_t)I)

  • Формальный вывод этого факта

xt=αtxt1+1αtzt1; где zt1,zt2,N(0,I)=αt(αt1xt2+1αt1zt2)+1αtzt1=αtαt1xt2+1αtαt1zˉt2; где zˉt2N(0,I)  ()==αˉtx0+1αˉtzq(xtx0)=N(xt;αˉtx0,(1αˉt)I)\begin{aligned} \mathbf{x}_t &= \sqrt{\alpha_t}\mathbf{x}_{t-1} + \sqrt{1 - \alpha_t}\mathbf{z}_{t-1}; \text{ где } \mathbf{z}_{t-1}, \mathbf{z}_{t-2}, \dots \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) \\ &= \sqrt{\alpha_t}(\sqrt{\alpha}_{t - 1}x_{t - 2} + \sqrt{1 - \alpha_{t - 1}}\mathbf{z}_{t - 2}) + \sqrt{1 - \alpha_t} \mathbf{z}_{t-1} \\ &= \sqrt{\alpha_t \alpha_{t-1}} \mathbf{x}_{t-2} + \sqrt{1 - \alpha_t \alpha_{t-1}} \bar{\mathbf{z}}_{t-2}; \text{ где } \bar{\mathbf{z}}_{t-2} \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) \ \ {(*)} \\ &= \dots \\ &= \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\mathbf{z} \\ q(\mathbf{x}_t \vert \mathbf{x}_0) &= \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t} \mathbf{x}_0, (1 - \bar{\alpha}_t)\mathbf{I}) \end{aligned}

(*) Пояснение ко второму переходу. У нас выходит

αt(1αt1)zt2+1αtzt1=αt(1αt1)+(1αt)zˉt2=1αtαt1zˉt2; где zt1,zt2,zˉt2N(0,I)\sqrt{\alpha_t(1 - \alpha_{t - 1})}z_{t - 2} + \sqrt{1 - \alpha_t}z_{t - 1} \\ = \sqrt{\alpha_t(1 - \alpha_{t - 1}) + (1 - \alpha_t)}\bar{z}_{t- 2} \\ = \sqrt{1 - \alpha_t\alpha_{t - 1}}\bar{z}_{t - 2}; \text{ где } z_{t - 1},z_{t - 2},\bar{z}_{t - 2} \sim \mathcal{N}(0, I)

Тогда LVLBL_{VLB} может быть переписано как

LVLB=Eq[DKL(q(xTx0)pθ(xT))LT++t=2TDKL(q(xt1xt,x0)pθ(xt1xt))Lt1logpθ(x0x1)L0L_{VLB} = \mathbb{E}_q [\underbrace{D_\text{KL}(q(\mathbf{x}_T \vert \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_T))}_{L_T} + + \sum_{t=2}^T \underbrace{D_\text{KL}(q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t))}_{L_{t-1}} \underbrace{- \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)}_{L_0}

  • Долгий вывод

Серым в скобках комментарий к последующему переходу.

LVLB=Eq(x0:T)[logq(x1:Tx0)pθ(x0:T)](расписываем совместное распределение)=Eq[logt=1Tq(xtxt1)pθ(xT)t=1Tpθ(xt1xt)]   (берем логарифм)=Eq[logpθ(xT)+t=1Tlogq(xtxt1)pθ(xt1xt)](отщепляем члены суммы)=Eq[logpθ(xT)+t=2Tlogq(xtxt1)pθ(xt1xt)+logq(x1x0)pθ(x0x1)](*)=Eq[logpθ(xT)+t=2Tlog(q(xt1xt,x0)pθ(xt1xt)q(xtx0)q(xt1x0))+logq(x1x0)pθ(x0x1)](лог произведения раскрываем)=Eq[logpθ(xT)+t=2Tlogq(xt1xt,x0)pθ(xt1xt)+t=2Tlogq(xtx0)q(xt1x0)+logq(x1x0)pθ(x0x1)](от второй суммы останется только 1ый и последний член)=Eq[logpθ(xT)+t=2Tlogq(xt1xt,x0)pθ(xt1xt)+logq(xTx0)q(x1x0)+logq(x1x0)pθ(x0x1)](комбинируем 1 и 3 член, 3 и 4 член)=Eq[logq(xTx0)pθ(xT)+t=2Tlogq(xt1xt,x0)pθ(xt1xt)logpθ(x0x1)]=Eq[DKL(q(xTx0)pθ(xT))LT+t=2TDKL(q(xt1xt,x0)pθ(xt1xt))Lt1logpθ(x0x1)L0]\begin{aligned}L_\text{VLB} &= \mathbb{E}_{q(\mathbf{x}_{0:T})} \Big[ \log\frac{q(\mathbf{x}_{1:T}\vert\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})} \Big] \quad \textit{\color{gray}{(расписываем совместное распределение)}}\\&= \mathbb{E}_q \Big[ \log\frac{\prod_{t=1}^T q(\mathbf{x}_t\vert\mathbf{x}_{t-1})}{ p_\theta(\mathbf{x}_T) \prod_{t=1}^T p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t) } \Big] \ \ \ \textit{\color{gray}{(берем логарифм)}} \\&= \mathbb{E}_q \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=1}^T \log \frac{q(\mathbf{x}_t\vert\mathbf{x}_{t-1})}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} \Big] \textit{\color{gray}{(отщепляем члены суммы)}} \\&= \mathbb{E}_q \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \frac{q(\mathbf{x}_t\vert\mathbf{x}_{t-1})}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \log\frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] \textit{\color{gray}{(*)}}\\&= \mathbb{E}_q \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \Big( \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)}\cdot \frac{q(\mathbf{x}_t \vert \mathbf{x}_0)}{q(\mathbf{x}_{t-1}\vert\mathbf{x}_0)} \Big) + \log \frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] \textit{\color{gray}{(лог произведения раскрываем)}}\\&= \mathbb{E}_q \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \sum_{t=2}^T \log \frac{q(\mathbf{x}_t \vert \mathbf{x}_0)}{q(\mathbf{x}_{t-1} \vert \mathbf{x}_0)} + \log\frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] \textit{\color{gray}{(от второй суммы останется только 1ый и последний член)}} \\&= \mathbb{E}_q \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \log\frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{q(\mathbf{x}_1 \vert \mathbf{x}_0)} + \log \frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] \textit{\color{gray}{(комбинируем 1 и 3 член, 3 и 4 член)}}\\&= \mathbb{E}_q \Big[ \log\frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_T)} + \sum_{t=2}^T \log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} - \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1) \Big] \\&= \mathbb{E}_q [\underbrace{D_\text{KL}(q(\mathbf{x}_T \vert \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_T))}_{L_T} + \sum_{t=2}^T \underbrace{D_\text{KL}(q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t))}_{L_{t-1}} \underbrace{- \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)}_{L_0} ]\end{aligned}

Пояснение (*). Пользуемся тем, что у нас Марковский процесс, и теоремой Байеса:

q(xtxt1)=q(xtxt1,x0)=q(xt1xt,x0)q(xtx0)q(xt1x0)q(x_t | x_{t - 1}) = q(x_t|x_{t - 1}, x_0) = \frac{q(x_{t - 1}| x_t, x_0)q(x_t | x_0)}{q(x_{t - 1}|x_0)}

Таким образом во время обучения, на каждой итерации параллельно оптимизируются случайные член LtL_t с помощью градиентного спуск (сэмлируем tU{1,...,T}t \sim U\{1,...,T\}). Поскольку все распределения нормальные, то KL между ними можно выписать в явной форме (см. ниже).

  • Формула KL между двумя нормальными

KL(N1  N2)=12(Tr(Σ21Σ1)+(μ2μ1)TΣ21(μ2μ1)++lndet(Σ2)det(Σ1)d)\begin{split} KL(\mathcal{N}_1 \ || \ \mathcal{N}_2) = \frac{1}{2}\bigg(Tr(\Sigma^{-1}_2 \Sigma_1) + (\mu_2 - \mu_1)^T \Sigma^{-1}_2 (\mu_2 - \mu_1) + \\[1.5ex] + \ln \frac{det(\Sigma_2)}{det(\Sigma_1)} - d \bigg) \end{split}

Если Σ1=σ1I, Σ2=σ2I\Sigma_1 = \sigma_1I, \ \Sigma_2 = \sigma_2I

KL(N1  N2)=12(dσ1σ2+1σ2μ2μ12++lnσ2σ1)\begin{split} KL(\mathcal{N}_1 \ || \ \mathcal{N}_2) = \frac{1}{2}\bigg(\frac{d\sigma_1}{\sigma_2} + \frac{1}{\sigma_2}\|\mu_2 - \mu_1\|^2 + \\[1.5ex] + \ln \frac{\sigma_2}{\sigma_1}\bigg) \end{split}

Осталось только выписать q(xt1xt,x0)q(x_{t - 1}| x_t, x_0) . Мы знаем, поскольку у нас все распределения нормальные, то и q(xt1xt,x0)q(x_{t - 1}| x_t, x_0) будет нормальным.

Обозначим

q(xt1xt,x0)=N(xt1;μ~(xt,x0),β~tI)q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_{t-1}; \color{blue}{\tilde{\boldsymbol{\mu}}}(\mathbf{x}_t, \mathbf{x}_0), \color{red}{\tilde{\beta}_t} \mathbf{I})

  • Вывод q(xt1xt,x0)q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)

Применим формулу Байеса и распишем. Тут мы просто пытаемся понять, как будут выглядеть среднее и дисперсия, выделяя квадратичную форму в показателе экспоненты

q(xt1xt,x0)=q(xtxt1,x0)q(xt1x0)q(xtx0)q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) = q(\mathbf{x}_t \vert \mathbf{x}_{t-1}, \mathbf{x}_0) \frac{ q(\mathbf{x}_{t-1} \vert \mathbf{x}_0) }{ q(\mathbf{x}_t \vert \mathbf{x}_0) } \propto

exp(12((xtαtxt1)2βt+(xt1αˉt1x0)21αˉt1(xtαˉtx0)21αˉt)) \propto \exp \Big(-\frac{1}{2} \big(\frac{(\mathbf{x}_t - \sqrt{\alpha_t} \mathbf{x}_{t-1})^2}{\beta_t} + \frac{(\mathbf{x}_{t-1} - \sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_0)^2}{1-\bar{\alpha}_{t-1}} - \frac{(\mathbf{x}_t - \sqrt{\bar{\alpha}_t} \mathbf{x}_0)^2}{1-\bar{\alpha}_t} \big) \Big)

=exp(12(xt22αtxtxt1+αtxt12βt+xt122αˉt1x0xt1+αˉt1x021αˉt1(xtαˉtx0)21αˉt))= = \exp \Big(-\frac{1}{2} \big(\frac{\mathbf{x}_t^2 - 2\sqrt{\alpha_t} \mathbf{x}_t \color{blue}{\mathbf{x}_{t-1}} \color{black}{+ \alpha_t} \color{red}{\mathbf{x}_{t-1}^2} }{\beta_t} + \frac{ \color{red}{\mathbf{x}_{t-1}^2} \color{black}{- 2 \sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_0} \color{blue}{\mathbf{x}_{t-1}} \color{black}{+ \bar{\alpha}_{t-1} \mathbf{x}_0^2} }{1-\bar{\alpha}_{t-1}} - \frac{(\mathbf{x}_t - \sqrt{\bar{\alpha}_t} \mathbf{x}_0)^2}{1-\bar{\alpha}_t} \big) \Big)=

=exp(12((αtβt+11αˉt1)xt12(2αtβtxt+2αˉt11αˉt1x0)xt1+C(xt,x0))) = \exp\Big( -\frac{1}{2} \big( \color{red}{(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}})} \mathbf{x}_{t-1}^2 - \color{blue}{(\frac{2\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t+\frac{2\sqrt{\bar{\alpha}_{t-1}}}{1 - \bar{\alpha}_{t-1}} \mathbf{x}_0)} \mathbf{x}_{t-1} \color{black}{ + C(\mathbf{x}_t, \mathbf{x}_0) \big) \Big)}

Далее перепишем красные и синие выражения в более красивой форме

β~t=1/(αtβt+11αˉt1)=1/(αtαˉt+βtβt(1αˉt1))=1αˉt11αˉtβt\color{red}{\tilde{\beta}_t} = 1/(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}}) = 1/(\frac{\alpha_t - \bar{\alpha}_t + \beta_t}{\beta_t(1 - \bar{\alpha}_{t-1})}) = \boxed{\frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t}

μ~t(xt,x0)=(αtβtxt+αˉt11αˉt1x0)/(αtβt+11αˉt1)\color{blue}{\tilde{\boldsymbol{\mu}}_t (\mathbf{x}_t, \mathbf{x}_0)} = (\frac{\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1} }}{1 - \bar{\alpha}_{t-1}} \mathbf{x}_0)/(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}})

=(αtβtxt+αˉt11αˉt1x0)1αˉt11αˉtβt = (\frac{\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1} }}{1 - \bar{\alpha}_{t-1}} \mathbf{x}_0)\frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t

=αt(1αˉt1)1αˉtxt+αˉt1βt1αˉtx0 = \boxed{\frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1 - \bar{\alpha}_t} \mathbf{x}_0}

Другой лосс. Предсказываем шум

В прошлой подсекции наша модель предсказывала среднее и дисперсию нормального распределения. Давайте зафиксируем Σθ(xt,t)=σt2I\Sigma_{\theta}(x_t, t) = \sigma^2_tI. Обычно берут σt2=βt\sigma^2_t = \beta_t или σt2=β~t=1αˉt11αˉtβt.\sigma^2_t = \tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t - 1}}{1 - \bar{\alpha}_t}\beta_t. Тогда Lt1L_{t - 1} из предыдущей секции можно переписать как

Lt=Eq[12σt2μθ(xt,x0)μ~t(xt,x0)2]+const(θ)L_{t} = \mathbb{E}_q\bigg[\frac{1}{2\sigma^2_t}\|\mu_{\theta}(x_t, x_0) - \tilde{\mu}_t(x_t, x_0)\|^2\bigg] + const(\theta)

Это первый момент, как меняется функционал, если мы не хотим предсказывать Σθ(xt,t)\Sigma_{\theta}(x_t, t), а фиксируем её.

Теперь вспомним, что q(xtx0)=N(xt;αˉtx0,(1αˉt)I)q(x_t | x_0) = \mathcal{N}(x_t;\sqrt{\bar{\alpha}_t} x_0, (1-\bar{\alpha}_t)I), но благодаря тому, что у нас гауссовское распределение, это можно переписать в виде

xt(x0,ϵ)=αˉtx0+1αˉtϵ,   ϵN(0,I)x_t(x_0, \epsilon) = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon, \ \ \ \epsilon \sim \mathcal{N}(0, I)

Выразим отсюда x0x_0 и получим, что x0=1αˉt(xt1αˉtϵ)x_0 = \frac{1}{\sqrt{\bar{\alpha}_t}}(x_t - \sqrt{1 - \bar{\alpha}_t}\epsilon), тогда подставим это выражение в формулу для μ~(xt,x0)\tilde{\boldsymbol{\mu}}(\mathbf{x}_t, \mathbf{x}_0) (из подсекции «Вывод q(xt1xt,x0)q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)») и получим

μ~(xt,x0)=1αt(xtβt1αˉtϵ)\tilde{\boldsymbol{\mu}}(\mathbf{x}_t, \mathbf{x}_0) = \frac{1}{\sqrt{\alpha_t}}(x_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}}\epsilon)

Теперь скажем, что наша модель будет предсказывать ϵ\epsilon. И просто будем «подставлять» его в выражение для μ~\tilde{\mu} выше. Обозначим предсказание модели как ϵθ(xt,t)\epsilon_{\theta}(x_t, t) — предсказанный шум ϵ\epsilon. Тогда лосс LtL_t превратиться в

Lt=Ex0,z[βt22αt(1αˉt)σt2ϵϵθ(αˉtx0+(1αˉt)ϵ,t)2]L_t = \mathbb{E}_{x_0, z}\bigg[ \frac{\beta^2_t}{2\alpha_t(1 - \bar{\alpha}_t )\sigma^2_t}\|\epsilon - \epsilon_{\theta}(\sqrt{\bar{\alpha}_t} x_0 + (1-\bar{\alpha}_t)\epsilon, t)\|^2\bigg]

Тем не менее лосс можно еще больше упростить и просто обучать с помощью MSE на ϵ\epsilon.

Ltsimple=Ex0,ϵ,t[ϵϵθ(αˉtx0+(1αˉt)ϵ,t)2]L^{simple}_t = \mathbb{E}_{x_0, \epsilon, t}\bigg[ \|\epsilon - \epsilon_{\theta}(\sqrt{\bar{\alpha}_t} x_0 + (1-\bar{\alpha}_t)\epsilon, t)\|^2\bigg]

Итак, алгоритмы обучения и сэмплирования выглядят вот так (на картинке z:=ϵz:= \epsilon).

Алгоритм обучения и сэмплирования диффузионной модели (Изображение взято из: Ho et al. 2020)

Стоит отметить, что важным недостатком диффузионных моделей является низкая скорость сэмплирования. Согласно Song et al. 2020: «Требуется 20 часов на генерацию 50 тысяч картинок размера 32х32, используя DDPM, и меньше минуты, используя GAN» (Nvidia 2080 Ti GPU). Тем не менее, в данном направлении был достигнут значительный прогресс и в целом проблема медленного сэмплирования была частично решена: Jiaming Song et al. (2021), Kong & Ping (2021), Bond-Taylor et al. (2021)

Давайте зафиксируем, какие функции потерь можно использовать. Для всех них справедлив тот факт, что мы сэмплируем шаг равномерно во время обучение tU{1,...,T})t \sim U\{1,...,T\}) и оптимизируем соответствующий LtL_{t}.

  • Оптимизируя член из суммы LVLBL_{VLB}. Это KL дивергенция между двумя нормальными распределениями

LVLB=EqDKL(q(xTx0)pθ(xT))LT++t=2TDKL(q(xt1xt,x0)pθ(xt1xt))Lt1logpθ(x0x1)L0L_{VLB} = \mathbb{E}_q \underbrace{D_\text{KL}(q(\mathbf{x}_T \vert \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_T))}_{L_T} + + \sum_{t=2}^T \underbrace{D_\text{KL}(q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t))}_{L_{t-1}} \underbrace{- \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)}_{L_0}

  • При фиксированной дисперсии Σθ\Sigma_{\theta} можно оптимизировать взвешенную MSE между средними нормальных распределений

Lt=Eq[12σt2μθ(xt,x0)μ~t(xt,x0)2]+const(θ)L_{t} = \mathbb{E}_q\bigg[\frac{1}{2\sigma^2_t}\|\mu_{\theta}(x_t, x_0) - \tilde{\mu}_t(x_t, x_0)\|^2\bigg] + const(\theta)

  • При фиксированной дисперсии и при предсказании шума с помощью взвешенной MSE. Или просто MSE. LtsimpleL^{simple}_t является самым популярным вариантом, который на практике дает лучшие результаты.

Lt=Ex0,z[βt22αt(1αˉt)σt2ϵϵθ(αˉtx0+(1αˉt)ϵ,t)2]Ltsimple=Ex0,z[ϵϵθ(αˉtx0+(1αˉt)ϵ,t)2]L_t = \mathbb{E}_{x_0, z}\bigg[ \frac{\beta^2_t}{2\alpha_t(1 - \bar{\alpha}_t )\sigma^2_t}\|\epsilon - \epsilon_{\theta}(\sqrt{\bar{\alpha}_t} x_0 + (1-\bar{\alpha}_t)\epsilon, t)\|^2\bigg] \\ L^{simple}_t = \mathbb{E}_{x_0, z}\bigg[ \|\epsilon - \epsilon_{\theta}(\sqrt{\bar{\alpha}_t} x_0 + (1-\bar{\alpha}_t)\epsilon, t)\|^2\bigg]

Выбор расписания βt\beta_t

Расписание является гиперпараметром, основными требованиями на который являются невозрастание (β1...βT)(\beta_1 \leq ... \leq \beta_T) и чтобы прямой процесс сходился к N(0,I)\mathcal{N}(0, I) в пределе по TT. Второе может гарантироваться тем, что αˉt0\bar{\alpha}_t \to 0. Вспомним,

q(xtx0)=N(xt;αˉtx0,(1αˉt)I)q(x_t | x_0) = \mathcal{N}(x_t;\sqrt{\bar{\alpha}_t} x_0, (1-\bar{\alpha}_t)I)

Однако на практике оно также проверяется, чтобы DKL(q(xTx0)N(0,I))D_{KL}(q(x_T | x_0) || \mathcal{N}(0, I)) было близко к 0.

Также стоит упомянуть, что обычно берут T=1000T = 1000. Но также важно помнить про требования выше, ведь расписание шума непосредственно зависит от TT.

Чаще всего используют линейное расписание, где β1=104, βT=0.02\beta_1 = 10^{-4}, \ \beta_T = 0.02. У данных констант нет никакой мотивации, кроме той, которая описана выше. Они были предложены в Ho et al. (2020).

В Nichol & Dhariwal (2021) было предложено косинусное расписание, которое помогло диффузионным моделям достичь лучшего NLL (negative loglikelihood):

βt=clip(1αˉtαˉt1,0.999)αˉt=f(t)f(0)where f(t)=cos(t/T+s1+sπ2)\beta_t = \text{clip}(1-\frac{\bar{\alpha}_t}{\bar{\alpha}_{t-1}}, 0.999) \quad\bar{\alpha}_t = \frac{f(t)}{f(0)}\quad\text{where }f(t)=\cos\Big(\frac{t/T+s}{1+s}\cdot\frac{\pi}{2}\Big)

Авторы обнаружили, что линейное расписание плохо работает на картинках 64х64 и меньше. А именно, последнии шаги прямого прохода были шумными и малоинформатиыными (просто зашумляем шум еще больше):

Пример зашумления картинки для линейного (сверху) и косинусного (снизу) расписания.

Также они обнаружили, что если обучать модель с линейным расписанием только на 80% первых шагов, то модель не становится сильно хуже, что подтверждает неиформативность последних шагов. Далее, они подобрали расписание так, чтобы αˉt\sqrt{\bar{\alpha}_t} убывало линейно на большей части отрезка (от 0 до TT) и почти не менялось рядом с 0 и TT. Разницу в αˉt\sqrt{\bar{\alpha}}_t для разных расписаний можно увидеть на картинке ниже:

Изображение взято из Nichol & Dhariwal, 2021

  • Детали

Также они ограничивают βt\beta_t числом 0.999, чтобы в конце процесса не было проблем с численной устойчивостью. Коэффициент ss используется, чтобы βt\beta_t не были слишком малы рядом с нулем. Он равен 0.008. Такое число было выбрано так, чтобы «β0\sqrt{\beta_0} была немного меньше, чем размер бина одного пикселя, то есть 1/127.51/127.5»

Classifier guidance

В Nichol & Dhariwal (2021) был предложен метод условной генерации, который повышает качество генерируемых картинок, при этом уменьшая их разнообразие. Для этого предобучается «шумный» классификатор на зашумленных картинках, то есть pϕ(yxt)p_{\phi}(y |x_t). Затем он используется во время сэмплирования, корректируя предсказанное среднее на xlogpϕ(yxt)\nabla_x \log p_{\phi}(y|x_t). В Nichol & Dhariwal (2021) (Секция 4.1) показывают, что данная добавка позволяет превратить распределение pθ(xi1xi)p_\theta(x_{i-1}|x_{i}) в pθ(xi1xi,y)p_\theta(x_{i-1}|x_{i}, y). Важно, что исходная диффузионная модель никак не меняется, что делает трюк еще более привлекательным. Алгоритм сэмплирования можно видеть на картинке ниже. Коэффициент ss отвечает за силу guidance.

  • Мотивация
    У генеративной модели GAN есть способ, который позволяет «балансировать» между разнообразием картинок и их качеством — truncation trick. Он заключается в сэмплировании латентного вектора truncated normal distibution. Данный трюк был хорошо описан и исследован в статье про BigGAN. Поэтому в диффузионных моделях тоже хотелось бы иметь метод, который позволяет балансировать между качеством и разнообразием. Авторы предложили classifier guidance, сравнили его с truncation trick и показали, что их метод строго лучше.

Изображение взято из Nichol & Dhariwal, 2021

Classifier-free guidance

Ho & Salimans (2021) предложили метод, в котором guidance достигается без использования дополнительной модели, поскольку это достаточно затратно. Для этого они обучали условную модель ϵθ(xty)\epsilon_{\theta}(x_t | y), у которой во время обучения реальная метка yy заменялась с какой-то фиксированной вероятностью (10%) на пустую метку (y=y=\emptyset). Это по сути позволяет нам обучать безусловную модель ϵθ(xt)\epsilon_{\theta}(x_t) одновременно с условной ϵθ(xty)\epsilon_{\theta}(x_t | y)Тогда во время сэмплирования делаем так, чтобы предсказание немного менялось в сторону ϵθ(xty)\epsilon_{\theta}(x_t | y), а именно:

ϵ^θ(xty)=ϵθ(xt)+s(ϵθ(xty)ϵθ(xt))\hat{\epsilon}_{\theta}(x_t | y) = \epsilon_{\theta}(x_t | \emptyset) + s\cdot(\epsilon_{\theta}(x_t | y) - \epsilon_{\theta}(x_t | \emptyset))

Мотивация этой формулы следовала из формулы Байеса:

p(yxt)p(xty)p(xt)    logp(yxt)logp(xty)logp(xt)    xtlogp(yxt)xtlogp(xty)xtlogp(xt)    xtlogp(yxt)ϵ(xty)ϵ(xt)p(y | x_t) \propto \frac{p(x_t | y) }{p(x_t)} \\ \implies \log p(y | x_t) \propto \log p(x_t | y) - \log p(x_t) \\ \implies \nabla_{x_t} \log p(y | x_t) \propto \nabla_{x_t} \log p(x_t | y) - \nabla_{x_t} \log p(x_t) \\ \implies \nabla_{x_t} \log p(y | x_t) \propto \epsilon(x_t | y) - \epsilon(x_t)

Тогда мы можем просто подставить xtlogp(yxt)\nabla_{x_t} \log p(y | x_t) в формулу для classifier guidance из предыдущей подсекции и получить желаемое равенство с точностью до коэффициента ss.

Овервью ключевых работ на сегодняшний день

Основная работа, в которой диффузионные модели (Denoising Diffusion Probabilistic Models, DDPMs) были применены для генерации картинок. Параграф в основном построен на ней.

Одна из первых попыток ускорить генерацию объектов. Идея следущая: давайте изменим прямой диффузионный процесс так, чтобы используя предобученную DDPM, приближать новый обратный процесс за меньшее число шагов.

Чтобы не обучать новую модель, нам нужен прямой диффузионный процесс, у которого будет такая же (суррогатная) функция потерь, а обратный процесс все еще останется Марковским. Оказалось, что существует целое семейство не-Марковских прямых процессов, удовлетворяющих этим требования. Это семейство имеет следующий вид:

qσ(x1:Tx0):=qσ(xTx0)t=2Tqσ(xt1xt,x0),q_\sigma (x_{1:T}|x_0):= q_\sigma (x_{T}|x_0)\prod_{t=2}^{T} q_\sigma (x_{t-1}|x_t,x_0),

где qσ(xTx0)=N(αtx0,(1αt)I)q_\sigma (x_{T}|x_0)= \mathcal{N}(\sqrt{\alpha_t}x_0, (1 - \alpha_t)I) и для всех t>1,t>1,

qσ(xt1xt,x0)=N(αt1x0+1αt1σt2xtαtx01αt,σt2I)q_\sigma (x_{t-1}|x_t,x_0)= \mathcal{N}(\sqrt{\alpha_{t-1}}x_0+\sqrt{1-\alpha_{t-1}-\sigma_t^2 }\cdot\frac{x_t-\sqrt{\alpha_{t}}x_0}{\sqrt{1-\alpha_t}},\sigma_t^2 I)

Среднее было выбрано так, чтобы qσ(xtx0)=N(αtx0,(1αt)I)q_{\sigma}(x_t | x_0) = \mathcal{N}(\sqrt{\alpha_t}x_0, (1 - \alpha_t)I) для всех tt. (см. Лемму 1 в Приложении B к статье). То есть важно лишь то, чтобы маргинальное распределение qσ(xtx0)q_{\sigma}(x_t | x_0) не менялось по сравнению с обычным Марковским случаем. Прямой процесс может быть получен с помощью теоремы Байеса:

qσ(xtxt1,x0)=qσ(xt1xt,x0)qσ(xtx0)qσ(xt1x0)q_\sigma (x_t|x_{t-1},x_0)=\frac{q_\sigma (x_{t-1}|x_t,x_0)q_\sigma (x_t|x_0)}{q_\sigma (x_{t-1}|x_0)}

Тут σ\sigma контролирует степень стохастичности прямого процесса. Можно заметить, что в отличии от исходного диффузионного процесса, предложенный прямой процесс больше не является Марковским, так как каждый xtx_t теперь зависит и от xt1x_{t-1} и от x0x_0. Схематично, это можно изобразить как на картинке справа. (Слева исходный диффузионный процесс для сравнения)

Screen

  • Заметка
    Авторы обращают внимание, что функция потерь в DDPM зависит от q(xtx0)q(x_t|x_0), а не от q(x1:xTx0)q(x_1{:}x_T | x_0) напрямую. Это означает, что нам нужно выбрать любой другой прямой диффузионный процесс, у которого q(xtx0)q(x_t|x_0) остались те же.

Далее, мы можем переписать обратный процесс в данном виде:

xt1=αt1  xt1αtϵθ(t)(xt)αt"predictedx0"+1αt1σt2ϵθ(t)(xt)"direction  pointing  to  xt"+σtϵtrandomnoisex_{t-1}=\sqrt{\alpha_{t-1}}\;\underbrace{\frac{x_t-\sqrt{1-\alpha_{t}}\epsilon_\theta^{(t)}(x_t)}{\sqrt{\alpha_t}}}_{"predicted\:x_0"}+\underbrace{\sqrt{1-\alpha_{t-1}-\sigma_t^2 }\cdot \epsilon_\theta^{(t)}(x_t)}_{"direction\;pointing\;to\;x_t"}+\underbrace{\sigma_t\epsilon_t}_{random\:noise}

Заметим, что при σt=(1αt1)(1αt)1αt/αt1\sigma_t = \sqrt{(1 - \alpha_{t - 1})(1 - \alpha_t)}\sqrt{1 - \alpha_t / \alpha_{t - 1}} прямой процесс становится марковским, а обратный как у DDPM (обычное сэмплирование, описанное в основной секции). При σt=0\sigma_t = 0 процесс сэмплирования становится детерминистичным (данный способ и называется DDIM). Ускорение сэмплирования достигается засчет использования лишь какого-то подмножества шагов (0τ1...τST,   S<T0 \leq \tau_1 \leq ... \leq \tau_S \leq T, \ \ \ S < T). Также одним из плюсов детерминистичного сэмплирования является возможность делать семантическую интерполяцию в латентном пространстве (как у GANов).

Улучшение DDPM, в котором был предложен новое расписание шума, что улучшило NLL. Также был изучен вариант, в котором дисперсия Σθ(xt,t)\Sigma_{\theta}(x_t, t) предсказывается моделью.

Статья, в которой показывается, что DDPM могут генерировать более качественные картинки по сравнению с GANами. Также был предложен метод conditional сэмплирования. Для этого предобучается классификатор на зашумленных сэмплах, а во время сэмплирования среднее нормального распределения «корректируется» на градиент классификатора.

Диффузионные модели на дискретных данных (например, текст). Вместо нормальных распределений используются категориальные. Также была обобщена мультиномиальная диффузия с помощью «матриц перехода», которые задают способ зашумления дискретных данных.

Более подробно: у нас есть xt{1,...,K}x_t \in \{1, ..., K\} — дискретная величина на всех шагах диффузии, тогда для каждого шага tt определена матрица прямого перехода QtQ_t такая, что [Qt]ij=q(xt=jxt1=i)[Q_t]_{ij} = q(x_t = j| x_{t - 1} = i). То есть строки матрицы суммируются в единицу. Тогда если обозначить через xtRK\mathbf{x}_t \in \mathbb{R}^K one-hot-закодированную версию xtx_t, то прямой процесс можно описать через категориальные распределения:

q(xtxt1)=Cat(xt;p=xt1Qt)q(\mathbf{x}_t | \mathbf{x}_{t -1}) = Cat(\mathbf{x}_t; \mathbf{p} = \mathbf{x}_{t - 1}Q_t)

Как и в нормальных распределениях, можем выписать

q(xtx0)=Cat(xt;p=x0Qˉt),   где   Qˉt=Q1Q2...Qtq(\mathbf{x}_t | \mathbf{x}_{0}) = Cat(\mathbf{x}_t; \mathbf{p} = \mathbf{x}_{0}\bar{Q}_t), \ \ \ где \ \ \ \bar{Q}_t=Q_1Q_2...Q_t

q(xt1xt,x0)=Cat(xt1;p=xtQtTx0Qˉt1x0QˉtxtT)q(\mathbf{x}_{t - 1}| \mathbf{x}_t, \mathbf{x}_0) = Cat(\mathbf{x}_{t - 1}; \mathbf{p} = \frac{\mathbf{x}_{t}{Q}^T_t \odot \mathbf{x}_0\bar{Q}_{t-1}}{\mathbf{x}_0\bar{Q}_t\mathbf{x}^T_t})

Поскольку тут нет такой хорошей параметризации через ϵ\epsilon, как у нормальных распределений, то единственный способ обучать — с помощью KL дивергенции (членамиLVLBL_{VLB}).

Остается только понять, как выбирать QtQ_t. Помимо того, чтобы сумма в каждой строчке была один, требуется, чтобы Qˉt\bar{Q}_t сходилось (при tt \to \infty) к равномерному распределению в каждой строчке (аналог нормального шума). За конкретными примерами стоит обратиться к статье.

Опишем работу метода GLIDE. Стоит задача генерировать картинки по заданному текстовому описанию. Для этого используется classifier-free guided diffusion model или CLIP. Это два разных варианта модели, которые авторы сравнивают. В первом случае модель обуславливается на эмбеддинги текста, которые были получены из обучаемого трансформера. Во втором случае guidance осуществляется за счет xtf(xt),g(c)\nabla_{x_t} \langle f(x_t), g(c) \rangle (это по сути градиент лосса метода CLIP) . Тут ff — это картиночный энкодер (на зашумленных картинках), а gg — это энкодер текстового входа. В целом, авторы получили, что classifier-free guidance генерирует более качественные картинки.

Способ описать диффузионные модели через стохастические дифференциальные уравнения.

Чтобы добавить в заметки выделенный текст, нажмите Command + E

Пройдите квиз по параграфу

Чтобы закрепить пройденный материал
Предыдущий параграф8.4. Нормализующие потоки
Следующий параграф8.6. Языковые модели

Вступайте в сообщество хендбука

Здесь можно найти единомышленников, экспертов и просто интересных собеседников. А ещё — получить помощь или поделиться знаниями.